#pragma OPENCL EXTENSION cl_khr_fp16 : enable
#define UP_DIV(x, y) (((x) + (y) - (1)) / (y))

__constant sampler_t
smp_none = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_NONE | CLK_FILTER_NEAREST;

constant FLT At[24] = {
        1.0000000000f, 1.0000000000f, 1.0000000000f, 1.0000000000f, 1.0000000000f, 0.0000000000f,
        0.0000000000f, 0.7071067691f, -0.7071067691f, 1.4142135382f, -1.4142135382f, 0.0000000000f,
        0.0000000000f, 0.4999999702f, 0.4999999702f, 1.9999998808f, 1.9999998808f, 0.0000000000f,
        0.0000000000f, 0.3535533845f, -0.3535533845f, 2.8284270763f, -2.8284270763f, 1.0000000000f
};

__kernel void Winograd36To4x4(__read_only image2d_t input,
                              __write_only image2d_t output,
                              __global FLT4 *bias,
                              int4 input_shape,      // N 36 TILE_XY CO_SLICES
                              int4 output_shape)     // N OH OW CO_SLICES
{
    int tile_xy = get_global_id(0);
    int row = get_global_id(1);
    int slice = get_global_id(2);

    int TILE_XY = input_shape.z;
    int SLICES = input_shape.w;
    int OH = output_shape.y;
    int OW = output_shape.z;

    if (tile_xy >= TILE_XY || row >= 4 || slice >= SLICES)
    {
        return;
    }

    constant FLT *At_row = At + row * 6;
    FLT4 AtM_row[6] = {0};
    for (int y = 0; y < 6; y++)
    {
        for (int x = 0; x < 6; x++)
        {
            AtM_row[x] += At_row[y] * READ_IMAGE(input, smp_none, (int2)(tile_xy, slice * 36 + y * 6 + x));
        }
    }

    int TILE_X = UP_DIV(OW, 4);
    for (int x = 0; x < 4; x++)
    {
        FLT4 acc = (FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
        for (int y = 0; y < 6; y++)
        {
            acc += AtM_row[y] * At[x * 6 + y];
        }
        acc += bias[slice];
        // acc = max(acc, (FLT4)(0.0f));

        int tile_x = tile_xy % TILE_X;
        int tile_y = tile_xy / TILE_X;
        int ow = tile_x * 4 + x;
        int oh = tile_y * 4 + row;
        if(oh==24 && ow==24 && slice==0)
        {
            printf("%.3f %.3f %.3f %.3f\n", acc.x, acc.y, acc.z, acc.w);
            printf("oh, ow,OH, OW: %d %d %d %d\n", oh, ow,OH, OW);
        }
        if (oh < OH)
            WRITE_IMAGE(output, (int2)(ow, slice * OH + oh), acc); // NC4HW4: height=CH width=W
//        if(ow < OW)
//            WRITE_IMAGE(output, (int2)(ow * SLICES + slice, oh), acc); // NHWC4 :height=H width=WC
    }
}